#include <string>
#include <iostream>
#include <fstream>
#include <sstream>

#include "clSetupAutomation.h"

namespace clAuto{
cl_platform_id platform_id;
cl_device_id device_id;
cl_uint num_platforms;
cl_uint num_devices;
cl_context context;
cl_command_queue command_queue;
int num_kernel;
cl_program *program = 0;
cl_kernel *kernel = 0;

#define clPrintf(DISPLAY, LOG) if (DISPLAY) \
	std::cout << LOG

static cl_program CreateProgram(cl_context aContext, cl_device_id aDevice, const char* aFileName)
{
	using namespace std;

	cl_int errNum;
	cl_program dProgram;

	ifstream srcFile(aFileName, ios::in);
	if (!srcFile.is_open())
	{
		cerr << "Failed to open file for reading: " << aFileName << endl;
		return NULL;
	}

	ostringstream oss;
	oss << srcFile.rdbuf();

	string srcProg = oss.str();
	const char *srcStr = srcProg.c_str();
	dProgram = clCreateProgramWithSource(context, 1/* the number of code */,
		(const char**)&srcStr, 
		NULL/* source and its size. the size is not given, null string represents exit character. */,
		&errNum/* error code */);
	if (dProgram == NULL)
	{
		cerr << "Failed to create CL program frome source." << endl;
		return NULL;
	}

	if(clBuildProgram(dProgram, 0,
		NULL/* specify building for device. NULL means building for every devices */,
		NULL/* build option */,
		NULL/* pfn_notify, which calls a notify routine when asynchronized building */,
		NULL/* arguments for pfn_notify */) != CL_SUCCESS)
	{
		cerr << "------------ Error in kernel: Log ---------------" << endl;
		char buildLog[16384];
		clGetProgramBuildInfo(dProgram, aDevice, CL_PROGRAM_BUILD_LOG, sizeof(buildLog), buildLog, NULL);
		cerr << buildLog << endl;
		cerr << "-------------------------------------------------" << endl;
		clReleaseProgram(dProgram);
		return NULL;
	}

	return dProgram;
}
static bool SaveProgramBinary(cl_program aProgram, cl_device_id aDevice, const char* aFileName)
{
	using namespace std;

	cl_uint numDevices = 0;
	cl_int errNum;
	bool complete = false;
	
	// Query for number of devices attached to program
	errNum = clGetProgramInfo(aProgram, CL_PROGRAM_NUM_DEVICES, sizeof(cl_uint), &numDevices, NULL);
	if (errNum != CL_SUCCESS)
	{
		cerr << "Error querying for number of devices." << endl;
		goto errCleanUp1;
	}

	// Get all of the Devices IDs
	cl_device_id *devices = new cl_device_id[numDevices];
	errNum = clGetProgramInfo(aProgram, CL_PROGRAM_DEVICES, sizeof(cl_device_id) * numDevices, devices, NULL);
	if (errNum != CL_SUCCESS)
	{
		cerr << "Error querying for devices." << endl;
		goto errCleanUp2;
	}

	// Determine the size of each program binary
	size_t *programBinarySizes = new size_t[numDevices];
	errNum = clGetProgramInfo(aProgram, CL_PROGRAM_BINARY_SIZES, sizeof(size_t) * numDevices, programBinarySizes, NULL);
	if (errNum != CL_SUCCESS)
	{
		cerr << "Error querying for program binary sizes." << endl;
		goto errCleanUp3;
	}

	unsigned char **programBinaries = new unsigned char*[numDevices];
	for (cl_uint i = 0; i < numDevices; i++)
		programBinaries[i] = new unsigned char[programBinarySizes[i]];

	// Get all of the program binaries
	errNum = clGetProgramInfo(aProgram, CL_PROGRAM_BINARIES, sizeof(unsigned char*) * numDevices, programBinaries, NULL);
	if (errNum != CL_SUCCESS)
	{
		cerr << "Error querying for program binaries" << endl;
		goto errCleanUp4;
	}

	// Finally store the binaries for the device requested out to disk for future reading.
	for (cl_uint i = 0; i < numDevices; i++)
	{
		if (devices[i] == aDevice)
		{
			FILE *fp = fopen(aFileName, "wb");
			fwrite(programBinaries[i], 1, programBinarySizes[i], fp);
			fclose(fp);
			break;
		}
	}

	// Cleanup
	complete = true;
errCleanUp4: 
	for (cl_uint i = 0; i < numDevices; i++)
		delete[] programBinaries[i];
	delete[] programBinaries;
errCleanUp3: delete[] programBinarySizes;
errCleanUp2: delete[] devices;
errCleanUp1: return complete;
}

static cl_program CreateProgramFromBinary(cl_context aContext, cl_device_id aDevice, const char* aFileName)
{
	using namespace std;

	FILE *fp = fopen(aFileName, "rb");
	if (fp == NULL)
		return NULL;

	// Determine the size of the binary
	size_t binarySize;
	fseek(fp, 0 ,SEEK_END);
	binarySize = ftell(fp);
	rewind(fp);
	
	// Load binary from disk
	unsigned char *programBinary = new unsigned char[binarySize];
	fread(programBinary, 1, binarySize, fp);
	fclose(fp);

	cl_int errNum = 0;
	cl_program dProgram;
	cl_int binaryStatus;

	dProgram = clCreateProgramWithBinary(aContext, 
		1, 
		&aDevice/* device id */, 
		&binarySize, 
		(const unsigned char**)&programBinary/* binary buffer */,
		&binaryStatus/* get binary status */,
		&errNum/* error num */);
	delete[] programBinary;
	if (errNum != CL_SUCCESS)
	{
		cerr << "Error loading program binary. " << endl;
		return NULL;
	}
	if (binaryStatus != CL_SUCCESS)
	{
		cerr << "Invalid binary for device." << endl;
		return NULL;
	}

	if(clBuildProgram(dProgram, 0, NULL, NULL, NULL, NULL) != CL_SUCCESS)
	{
		cerr << aFileName << ": Failed to create CL program from binary." << endl;
		cerr << "--------------------Build Log--------------------" << endl;
		char buildLog[16384];
		clGetProgramBuildInfo(dProgram, device_id, CL_PROGRAM_BUILD_LOG, sizeof(buildLog), buildLog, NULL);
		cerr << buildLog << endl;
		cerr << "-------------------------------------------------" << endl;
		clReleaseProgram(dProgram);
		return NULL;
	}
	return dProgram;
}

void clSetupAutomation(char **kernel_name, int nKernel, int loadKernelFromBinary, int dispLog, int default_device)
{
	using namespace std;
	stringstream tss;
	cl_int errNum;
	clGetPlatformIDs(1, &platform_id, &num_platforms);
	
	clPrintf(dispLog, "Get the Device info and select Device...\n");
	errNum = clGetDeviceIDs(platform_id,//the platform id
		CL_DEVICE_TYPE_GPU,//device_type
		0,//the number of cl_device entries
		NULL, &num_devices);
	cl_device_id * cdDevices = (cl_device_id*)malloc(sizeof(cl_device_id*) * num_devices);
	errNum |= clGetDeviceIDs(platform_id,	CL_DEVICE_TYPE_GPU, num_devices, cdDevices, NULL);
	tss << "   # of Devices Available = " << num_devices << endl;
	clPrintf(dispLog, tss.str());
	tss.str("");

	// Select Device and show device info
	cl_uint target_device;
	if (num_devices == 1)
		target_device = 1;
	else
		target_device = default_device;
	tss << "   Using Device " << target_device << ": ";
	clPrintf(dispLog, tss.str());
	tss.str("");
	device_id = cdDevices[target_device-1];
	free(cdDevices);
	char device_string[1024];
	errNum |=clGetDeviceInfo(device_id,//device id
		CL_DEVICE_NAME,//device information queried
		sizeof(device_string),//device information buffer size
		&device_string,//device information buffer
		NULL//actual size of device information
		);
	clPrintf(dispLog, device_string);
	cl_uint num_compute_units;
	errNum |=clGetDeviceInfo(device_id, CL_DEVICE_MAX_COMPUTE_UNITS, sizeof(num_compute_units), &num_compute_units, NULL);
	tss << endl << "   # of Compute Units = " << num_compute_units << endl;
	clPrintf(dispLog, tss.str());
	tss.str("");
	
	clPrintf(dispLog, "Make-up Context ...");
	context = clCreateContext(NULL,// properties, NULL means the platform is implementation-defined
		1,// the number of devices specified in next argument(_device_id)
		&device_id,// the device id
		NULL,//pfn_notify
		NULL,//user_data
		&errNum//error code
		);
	if(errNum != CL_SUCCESS) goto clean_up;

	clPrintf(dispLog, "OK! Allocate Command Queue ...");
	command_queue = clCreateCommandQueue(context, device_id,
		NULL,//CL_QUEUE_OUT_OF_ORDER_EXEC_MODE_ENABLE | CL_QUEUE_PROFILING_ENABLE
		&errNum//error code
		);
	if(errNum != CL_SUCCESS) goto clean_up;

	/// Compile a kernel program 
	num_kernel = nKernel;
	tss << "OK!\nThere are " << num_kernel << " kernels. Start loading and building them." << endl;
	clPrintf(dispLog, tss.str());
	tss.str("");
	program = (cl_program*)malloc(sizeof(cl_program)*num_kernel);
	kernel = (cl_kernel*)malloc(sizeof(cl_kernel)*num_kernel);
	for (int i = 0; i < num_kernel; i++)
	{
		program[i] = NULL;
		kernel[i] = NULL;

		stringstream sourceFileName;
		stringstream binaryFileName;
		sourceFileName << kernel_name[i] << ".cl";
		binaryFileName << sourceFileName.str() << ".bin";
		if (loadKernelFromBinary)
			program[i] = CreateProgramFromBinary(context, device_id, binaryFileName.str().c_str());
		else
			program[i] = NULL;

		if (program[i] == NULL)
		{// there isn't binary source
			program[i] = CreateProgram(context, device_id, sourceFileName.str().c_str());
			if (program[i] == NULL)
				goto clean_up;

			// Save the compiled binary
			if (SaveProgramBinary(program[i], device_id, binaryFileName.str().c_str()) == false)
			{
				cerr << kernel_name[i] << ": Failed to write program binary." << endl;
				goto clean_up;
			}
			cout << kernel_name[i] << ": Read program from source." << endl;
		}
		else // there is already binary source
			cout << kernel_name[i] << ": Read program from binary." << endl;

		kernel[i] = clCreateKernel(program[i], kernel_name[i], &errNum/* error code */);
		if (errNum != CL_SUCCESS) goto clean_up;
	}
	clPrintf(dispLog, "Program Build Success\n\n");

	return;

clean_up:
	clPrintf(dispLog, "FAIL! So, Clean-Up ...");
	clClearAutomation();
	clPrintf(dispLog, "Finished.\n\n");
	exit(-1);
}

void clClearAutomation(void)
{
	if (command_queue != NULL) 
	{
		clFlush(command_queue);
		clFinish(command_queue);
	}
	if (kernel != NULL)
	{
		for (int i = 0; i < num_kernel; i++)
		{
			if (kernel[i] != NULL)
			{
				clReleaseKernel(kernel[i]);
				kernel[i] = NULL;
			}
		}
		free(kernel);
		kernel = NULL;
	}
	if (program != NULL) 
	{
		for (int i = 0; i < num_kernel; i++)
		{
			if (program[i] != NULL)
			{
				clReleaseProgram(program[i]);
				program[i] = NULL;
			}
		}
		free(program);
		program = NULL;
	}
	if (command_queue != 0) clReleaseCommandQueue(command_queue);
	if (context != 0) clReleaseContext(context);
}
};